from __future__ import print_function, division
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.data
import torch

from SwinTransformer.models.swin_transformer import ViTSwinTransformerBlock, VitPatchMerging


class Partial_conv3(nn.Module):

    def __init__(self, dim, n_div, forward):
        super().__init__()
        self.dim_conv3 = dim // n_div
        self.dim_untouched = dim - self.dim_conv3
        self.partial_conv3 = nn.Conv2d(self.dim_conv3, self.dim_conv3, 3, 1, 1, bias=True)

        if forward == 'slicing':
            self.forward = self.forward_slicing
        elif forward == 'split_cat':
            self.forward = self.forward_split_cat
        else:
            raise NotImplementedError

    def forward_slicing(self, x):
        # only for inference
        x = x.clone()  # !!! Keep the original input intact for the residual connection later
        x[:, :self.dim_conv3, :, :] = self.partial_conv3(x[:, :self.dim_conv3, :, :])

        return x

    def forward_split_cat(self, x):
        # for training/inference
        x1, x2 = torch.split(x, [self.dim_conv3, self.dim_untouched], dim=1)
        x1 = self.partial_conv3(x1)
        x = torch.cat((x1, x2), 1)

        return x


class conv_block(nn.Module):
    """
    Convolution Block
    """

    def __init__(self, in_ch, out_ch):
        super(conv_block, self).__init__()

        self.conv = nn.Sequential(
            Partial_conv3(in_ch, 2, 'split_cat'),
            nn.Conv2d(in_ch, out_ch, kernel_size=3, stride=1, padding=1, bias=True),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True),

            # nn.Conv2d(out_ch, out_ch, kernel_size=3, stride=1, padding=1, bias=True),
            Partial_conv3(out_ch, 2, 'split_cat'),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True))

    def forward(self, x):
        x = self.conv(x)

        return x


class up_conv(nn.Module):
    """
    Up Convolution Block
    """

    def __init__(self, in_ch, out_ch, scale_factor=2):
        super(up_conv, self).__init__()
        self.up = nn.Sequential(
            nn.Upsample(scale_factor=scale_factor),
            nn.Conv2d(in_ch, out_ch, kernel_size=3, stride=1, padding=1, bias=True),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        x = self.up(x)
        return x


class ConvBNRelu1D(nn.Module):
    def __init__(self, inChannel, outChannel, kernel_size=3, padding=1, bias=False):
        super(ConvBNRelu1D, self).__init__()
        self.conv = nn.Conv1d(inChannel, outChannel, kernel_size=kernel_size, padding=padding, bias=bias)
        self.bn = nn.BatchNorm1d(outChannel, momentum=0.1, affine=True)
        self.reLu = nn.ReLU(inplace=True)

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        return self.reLu(x)


class LGFO(nn.Module):
    """
    UNet - Basic Implementation
    Paper : https://arxiv.org/abs/1505.04597
    """

    def __init__(self, in_ch, hide_ch, out_ch, inWidth, hideWidth, outWidth):
        super().__init__()

        # self.layer1 =  nn.Linear(inWidth * inWidth, hideWidth * hideWidth)
        # self.CBR1=    ConvBNRelu1D(in_ch, hide_ch, kernel_size=1, padding=0, bias=True)
        #
        # self.layer2 = nn.Linear(hideWidth * hideWidth ,hideWidth * hideWidth)
        # self.CBR2 =    ConvBNRelu1D(hide_ch, hide_ch, kernel_size=1, padding=0, bias=True)
        #
        # self.layer3= nn.Linear(hideWidth * hideWidth, outWidth * outWidth)
        # self.CBR3 =    ConvBNRelu1D(hide_ch, out_ch, kernel_size=1, padding=0, bias=True)

        self.layer1 = ConvBNRelu1D(in_ch, out_ch, kernel_size=1, padding=0, bias=True)
        self.layer2 = nn.Linear(inWidth * inWidth, hideWidth * hideWidth)
        # ConvBNRelu1D(hide_ch, hide_ch, kernel_size=1, padding=0, bias=True),
        self.layer3 = nn.Linear(hideWidth * hideWidth, hideWidth * hideWidth)
        self.layer4 = nn.Linear(hideWidth * hideWidth, inWidth * inWidth)
        self.layer5 = ConvBNRelu1D(out_ch, out_ch, kernel_size=1, padding=0, bias=True)

        self.outDim = out_ch

    def forward(self, x):
        # windows间交互
        b, c, w, h = x.shape
        x = x.view(b, c, w * h)

        x = self.layer1(x)
        shortcut = x
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        x = x + shortcut
        x = self.layer5(x)

        x = x.view(b, self.outDim, w, h)
        x = x.reshape(b, self.outDim, w, h)

        return x


class LGFO_PConv(nn.Module):
    """
    UNet - Basic Implementation
    Paper : https://arxiv.org/abs/1505.04597
    """

    def __init__(self, in_ch=3, out_ch=1):
        super().__init__()

        n1 = 64
        filters = [n1, n1 * 2, n1 * 4, n1 * 8, n1 * 16]

        self.Maxpool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.Maxpool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.Maxpool3 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.Maxpool4 = nn.MaxPool2d(kernel_size=2, stride=2)

        self.Conv1 = conv_block(in_ch, filters[0])

        self.Conv2 = conv_block(filters[0], filters[1])
        self.Conv3 = conv_block(filters[1], filters[2])
        self.Conv4 = conv_block(filters[2], filters[3])
        # self.Conv4=WinLearnedAttention(64,filters[2], filters[3])
        self.Conv5 = conv_block(filters[3], filters[4])

        # self.pConv1 = pConv.Partial_conv3(filters[0], 2, 'split_cat')
        # self.pConv2 = pConv.Partial_conv3(filters[1], 2, 'split_cat')
        # self.pConv3 = pConv.Partial_conv3(filters[2], 2, 'split_cat')
        # self.pConv4 = pConv.Partial_conv3(filters[3], 2, 'split_cat')
        # self.pConv5 = pConv.Partial_conv3(filters[4], 2, 'split_cat')

        # self.Conv5=WinLearnedAttention(32,filters[3], filters[4])

        self.Up5 = up_conv(filters[4], filters[3])
        self.Up_conv5 = conv_block(filters[4], filters[3])
        # self.Up_conv5 = WinLearnedAttention(64,filters[4], filters[3])

        self.Up4 = up_conv(filters[3], filters[2])
        self.Up_conv4 = conv_block(filters[3], filters[2])

        self.Up3 = up_conv(filters[2], filters[1])
        self.Up_conv3 = conv_block(filters[2], filters[1])

        self.Up2 = up_conv(filters[1], filters[0])
        self.Up_conv2 = conv_block(filters[1], filters[0])

        self.Conv = nn.Conv2d(filters[0], out_ch, kernel_size=1, stride=1, padding=0)

        # self.active = torch.nn.Sigmoid()
        self.interWin1 = LGFO(3, filters[0], filters[0], 256, 32, 256)
        self.interWin2 = LGFO(filters[0], filters[1], filters[1], 128, 32, 128)
        self.interWin3 = LGFO(filters[1], filters[2], filters[2], 64, 32, 64)
        self.interWin4 = LGFO(filters[2], filters[3], filters[3], 32, 32, 32)
        self.interWin5 = LGFO(filters[3], filters[4], filters[4], 16, 16, 16)

    def forward(self, input):
        g1 = self.interWin1(input)

        g2 = self.Maxpool1(g1)
        g2 = self.interWin2(g2)

        g3 = self.Maxpool1(g2)
        g3 = self.interWin3(g3)

        g4 = self.Maxpool1(g3)
        g4 = self.interWin4(g4)

        g5 = self.Maxpool1(g4)
        g5 = self.interWin5(g5)

        e1 = self.Conv1(input)
        # e1=self.pConv1(e1)

        e2 = self.Maxpool1(e1)
        e2 = self.Conv2(e2)
        # e2=self.pConv2(e2)

        e3 = self.Maxpool2(e2)
        e3 = self.Conv3(e3)
        # e3=self.pConv3(e3)

        e4 = self.Maxpool3(e3)
        e4 = self.Conv4(e4)
        # e4=self.pConv4(e4)

        e5 = self.Maxpool4(e4)
        e5 = self.Conv5(e5)
        # e5=self.pConv5(e5)

        e1 = e1 + g1
        e2 = e2 + g2
        e3 = e3 + g3
        e4 = e4 + g4
        e5 = e5 + g5

        d5 = self.Up5(e5)
        d5 = torch.cat((e4, d5), dim=1)

        d5 = self.Up_conv5(d5)

        d4 = self.Up4(d5)
        d4 = torch.cat((e3, d4), dim=1)
        d4 = self.Up_conv4(d4)

        d3 = self.Up3(d4)
        d3 = torch.cat((e2, d3), dim=1)
        d3 = self.Up_conv3(d3)

        d2 = self.Up2(d3)
        d2 = torch.cat((e1, d2), dim=1)
        d2 = self.Up_conv2(d2)

        out = self.Conv(d2)

        # d1 = self.active(out)

        return out


class LGFO_Net_NoConv(nn.Module):
    """
    UNet - Basic Implementation
    Paper : https://arxiv.org/abs/1505.04597
    """

    def __init__(self, in_ch=3, out_ch=1):
        super().__init__()

        n1 = 64
        filters = [n1, n1 * 2, n1 * 4, n1 * 8, n1 * 16]

        self.Maxpool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.Maxpool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.Maxpool3 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.Maxpool4 = nn.MaxPool2d(kernel_size=2, stride=2)

        self.interWin1 = LGFO(3, filters[0], filters[0], 256, 32, 256)
        self.interWin2 = LGFO(filters[0], filters[1], filters[1], 128, 32, 128)
        self.interWin3 = LGFO(filters[1], filters[2], filters[2], 64, 32, 64)
        self.interWin4 = LGFO(filters[2], filters[3], filters[3], 32, 32, 32)
        self.interWin5 = LGFO(filters[3], filters[4], filters[4], 16, 16, 16)

        self.UNetDecoder = UnetDecoder(depths=filters)

    def forward(self, input):
        g1 = self.interWin1(input)

        g2 = self.Maxpool1(g1)
        g2 = self.interWin2(g2)

        g3 = self.Maxpool1(g2)
        g3 = self.interWin3(g3)

        g4 = self.Maxpool1(g3)
        g4 = self.interWin4(g4)

        g5 = self.Maxpool1(g4)
        g5 = self.interWin5(g5)

        out = self.UNetDecoder([g1, g2, g3, g4, g5])

        return out


class mutiDecoder(nn.Module):
    def __init__(self, depths=[32, 32, 32, 32, 32]):
        super().__init__()

        self.ConvTranspose1_0 = up_conv(depths[1], depths[0])
        self.Conv1_0 = conv_block(depths[0], depths[0])

        self.ConvTranspose2_1 = up_conv(depths[2], depths[1])
        self.Conv2_1 = conv_block(depths[1] * 2, depths[1])

        self.ConvTranspose2_0 = up_conv(depths[1], depths[0])
        self.Conv2_0 = conv_block(depths[0], depths[0])

        self.ConvTranspose3_2 = up_conv(depths[3], depths[2])
        self.Conv3_2 = conv_block(depths[2] * 2, depths[2])

        self.ConvTranspose3_1 = up_conv(depths[2], depths[1])
        self.Conv3_1 = conv_block(depths[1] * 2, depths[1])

        self.ConvTranspose3_0 = up_conv(depths[1], depths[0])
        self.Conv3_0 = conv_block(depths[0], depths[0])

        self.ConvTranspose4_3 = up_conv(depths[4], depths[3])
        self.Conv4_3 = conv_block(depths[3] * 2, depths[3])

        self.ConvTranspose4_2 = up_conv(depths[3], depths[2])
        self.Conv4_2 = conv_block(depths[2] * 2, depths[2])

        self.ConvTranspose4_1 = up_conv(depths[2], depths[1])
        self.Conv4_1 = conv_block(depths[1] * 2, depths[1])

        self.ConvTranspose4_0 = up_conv(depths[1], depths[0])
        self.Conv4_0 = conv_block(depths[0], depths[0])

        self.ConvAll = conv_block(depths[0] * 5, depths[0])
        self.Out = nn.Conv2d(depths[0], 2, kernel_size=1, stride=1, padding=0)

    def forward(self, levels):
        level0 = levels[0]
        # level1的解码
        level1 = self.ConvTranspose1_0(levels[1])
        level1 = self.Conv1_0(level1)

        # level2的解码
        level2 = self.ConvTranspose2_1(levels[2])
        level2 = torch.cat([level2, levels[1]], dim=1)
        level2 = self.Conv2_1(level2)

        level2 = self.ConvTranspose2_0(level2)
        level2 = self.Conv2_0(level2)

        # level3的解码
        level3 = self.ConvTranspose3_2(levels[3])
        level3 = torch.cat([level3, levels[2]], dim=1)
        level3 = self.Conv3_2(level3)

        level3 = self.ConvTranspose3_1(level3)
        level3 = torch.cat([level3, levels[1]], dim=1)
        level3 = self.Conv3_1(level3)

        level3 = self.ConvTranspose3_0(level3)
        level3 = self.Conv3_0(level3)

        # level4的解码
        level4 = self.ConvTranspose4_3(levels[4])
        level4 = torch.cat([level4, levels[3]], dim=1)
        level4 = self.Conv4_3(level4)

        level4 = self.ConvTranspose4_2(level4)
        level4 = torch.cat([level4, levels[2]], dim=1)
        level4 = self.Conv4_2(level4)

        level4 = self.ConvTranspose4_1(level4)
        level4 = torch.cat([level4, levels[1]], dim=1)
        level4 = self.Conv4_1(level4)

        level4 = self.ConvTranspose4_0(level4)
        level4 = self.Conv4_0(level4)

        # 所有level的融合
        out = torch.cat([level0, level1, level2, level3, level4], dim=1)
        out = self.ConvAll(out)
        out = self.Out(out)

        return out


class UnetDecoder(nn.Module):
    def __init__(self, depths=[32, 32, 32, 32, 32]):
        super().__init__()

        self.ConvTranspose4_3 = up_conv(depths[4], depths[3])
        self.Conv4_3 = conv_block(depths[3] * 2, depths[3])

        self.ConvTranspose4_2 = up_conv(depths[3], depths[2])
        self.Conv4_2 = conv_block(depths[2] * 2, depths[2])

        self.ConvTranspose4_1 = up_conv(depths[2], depths[1])
        self.Conv4_1 = conv_block(depths[1] * 2, depths[1])

        self.ConvTranspose4_0 = up_conv(depths[1], depths[0])
        self.Conv4_0 = conv_block(depths[0], depths[0])

        self.Out = nn.Conv2d(depths[0], 2, kernel_size=1, stride=1, padding=0)

    def forward(self, levels):
        # level4的解码
        level4 = self.ConvTranspose4_3(levels[4])
        level4 = torch.cat([level4, levels[3]], dim=1)
        level4 = self.Conv4_3(level4)

        level4 = self.ConvTranspose4_2(level4)
        level4 = torch.cat([level4, levels[2]], dim=1)
        level4 = self.Conv4_2(level4)

        level4 = self.ConvTranspose4_1(level4)
        level4 = torch.cat([level4, levels[1]], dim=1)
        level4 = self.Conv4_1(level4)

        level4 = self.ConvTranspose4_0(level4)
        level4 = self.Conv4_0(level4)

        out = self.Out(level4)

        return out


class DenseDecoder(nn.Module):
    def __init__(self, depths=[32, 32, 32, 32, 32]):
        super().__init__()

        self.ConvTranspose1_0 = up_conv(depths[1], depths[0], scale_factor=2)

        self.ConvTranspose2_1 = up_conv(depths[2], depths[1], scale_factor=2)

        self.ConvTranspose2_0 = up_conv(depths[2], depths[0], scale_factor=4)

        self.ConvTranspose3_2 = up_conv(depths[3], depths[2], scale_factor=2)

        self.ConvTranspose3_1 = up_conv(depths[3], depths[1], scale_factor=4)

        self.ConvTranspose3_0 = up_conv(depths[3], depths[0], scale_factor=8)

        self.ConvTranspose4_3 = up_conv(depths[4], depths[3], scale_factor=2)

        self.ConvTranspose4_2 = up_conv(depths[4], depths[2], scale_factor=4)

        self.ConvTranspose4_1 = up_conv(depths[4], depths[1], scale_factor=8)

        self.ConvTranspose4_0 = up_conv(depths[4], depths[0], scale_factor=16)

        self.upConv3 = conv_block(depths[3] * 2, depths[3])
        self.upConv2 = conv_block(depths[2] * 3, depths[2])
        self.upConv1 = conv_block(depths[1] * 4, depths[1])
        self.upConv0 = conv_block(depths[0] * 5, depths[0])

        self.out = nn.Conv2d(depths[0], 2, kernel_size=1, stride=1, padding=0)

    def forward(self, levels):
        l4_3 = self.ConvTranspose4_3(levels[4])
        l4_2 = self.ConvTranspose4_2(levels[4])
        l4_1 = self.ConvTranspose4_1(levels[4])
        l4_0 = self.ConvTranspose4_0(levels[4])

        levels[3] = torch.cat([levels[3], l4_3], dim=1)
        levels[3] = self.upConv3(levels[3])

        l3_2 = self.ConvTranspose3_2(levels[3])
        l3_1 = self.ConvTranspose3_1(levels[3])
        l3_0 = self.ConvTranspose3_0(levels[3])

        levels[2] = torch.cat([levels[2], l4_2, l3_2], dim=1)
        levels[2] = self.upConv2(levels[2])

        l2_1 = self.ConvTranspose2_1(levels[2])
        l2_0 = self.ConvTranspose2_0(levels[2])

        levels[1] = torch.cat([levels[1], l4_1, l3_1, l2_1], dim=1)
        levels[1] = self.upConv1(levels[1])

        l1_0 = self.ConvTranspose1_0(levels[1])

        levels[0] = torch.cat([levels[0], l4_0, l3_0, l2_0, l1_0], dim=1)
        levels[0] = self.upConv0(levels[0])

        out = self.out(levels[0])

        return out


class LGFO_Net1(nn.Module):
    """
    UNet - Basic Implementation
    Paper : https://arxiv.org/abs/1505.04597
    """

    def __init__(self, in_ch=3, out_ch=1):
        super().__init__()

        n1 = 64
        filters = [n1, n1 * 2, n1 * 4, n1 * 8, n1 * 16]

        self.Maxpool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.Maxpool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.Maxpool3 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.Maxpool4 = nn.MaxPool2d(kernel_size=2, stride=2)

        self.Conv1 = conv_block(in_ch, filters[0])
        self.Conv2 = conv_block(filters[0], filters[1])
        self.Conv3 = conv_block(filters[1], filters[2])
        self.Conv4 = conv_block(filters[2], filters[3])
        # self.Conv4=WinLearnedAttention(64,filters[2], filters[3])
        self.Conv5 = conv_block(filters[3], filters[4])
        # self.Conv5=WinLearnedAttention(32,filters[3], filters[4])

        # self.active = torch.nn.Sigmoid()
        self.interWin1 = LGFO(3, filters[0], filters[0], 256, 32, 256)
        self.interWin2 = LGFO(filters[0], filters[1], filters[1], 128, 32, 128)
        self.interWin3 = LGFO(filters[1], filters[2], filters[2], 64, 32, 64)
        self.interWin4 = LGFO(filters[2], filters[3], filters[3], 32, 32, 32)
        self.interWin5 = LGFO(filters[3], filters[4], filters[4], 16, 16, 16)

        self.decoder = decoder(filters)

    def forward(self, input):
        g1 = self.interWin1(input)

        g2 = self.Maxpool1(g1)
        g2 = self.interWin2(g2)

        g3 = self.Maxpool1(g2)
        g3 = self.interWin3(g3)

        g4 = self.Maxpool1(g3)
        g4 = self.interWin4(g4)

        g5 = self.Maxpool1(g4)
        g5 = self.interWin5(g5)

        e1 = self.Conv1(input)

        e2 = self.Maxpool1(e1)
        e2 = self.Conv2(e2)

        e3 = self.Maxpool2(e2)
        e3 = self.Conv3(e3)

        e4 = self.Maxpool3(e3)
        e4 = self.Conv4(e4)

        e5 = self.Maxpool4(e4)
        e5 = self.Conv5(e5)

        e1 = e1 + g1
        e2 = e2 + g2
        e3 = e3 + g3
        e4 = e4 + g4
        e5 = e5 + g5

        out = self.decoder([e1, e2, e3, e4, e5])

        # d1 = self.active(out)

        return out


class LGFO_Net_denseDecode(nn.Module):
    """
    UNet - Basic Implementation
    Paper : https://arxiv.org/abs/1505.04597
    """

    def __init__(self, in_ch=3, out_ch=1):
        super().__init__()

        n1 = 64
        filters = [n1, n1 * 2, n1 * 4, n1 * 8, n1 * 16]

        self.Maxpool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.Maxpool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.Maxpool3 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.Maxpool4 = nn.MaxPool2d(kernel_size=2, stride=2)

        self.Conv1 = conv_block(in_ch, filters[0])
        self.Conv2 = conv_block(filters[0], filters[1])
        self.Conv3 = conv_block(filters[1], filters[2])
        self.Conv4 = conv_block(filters[2], filters[3])
        # self.Conv4=WinLearnedAttention(64,filters[2], filters[3])
        self.Conv5 = conv_block(filters[3], filters[4])
        # self.Conv5=WinLearnedAttention(32,filters[3], filters[4])

        # self.active = torch.nn.Sigmoid()
        self.interWin1 = LGFO(3, filters[0], filters[0], 256, 32, 256)
        self.interWin2 = LGFO(filters[0], filters[1], filters[1], 128, 32, 128)
        self.interWin3 = LGFO(filters[1], filters[2], filters[2], 64, 32, 64)
        self.interWin4 = LGFO(filters[2], filters[3], filters[3], 32, 32, 32)
        self.interWin5 = LGFO(filters[3], filters[4], filters[4], 16, 16, 16)

        self.decoder = DenseDecoder(filters)

    def forward(self, input):
        g1 = self.interWin1(input)

        g2 = self.Maxpool1(g1)
        g2 = self.interWin2(g2)

        g3 = self.Maxpool1(g2)
        g3 = self.interWin3(g3)

        g4 = self.Maxpool1(g3)
        g4 = self.interWin4(g4)

        g5 = self.Maxpool1(g4)
        g5 = self.interWin5(g5)

        e1 = self.Conv1(input)

        e2 = self.Maxpool1(e1)
        e2 = self.Conv2(e2)

        e3 = self.Maxpool2(e2)
        e3 = self.Conv3(e3)

        e4 = self.Maxpool3(e3)
        e4 = self.Conv4(e4)

        e5 = self.Maxpool4(e4)
        e5 = self.Conv5(e5)

        e1 = e1 + g1
        e2 = e2 + g2
        e3 = e3 + g3
        e4 = e4 + g4
        e5 = e5 + g5

        out = self.decoder([e1, e2, e3, e4, e5])

        # d1 = self.active(out)

        return out


# use swintransformerblock and LGFO architecture
class LGFO_SwinTransformerNet(nn.Module):
    """
    UNet - Basic Implementation
    Paper : https://arxiv.org/abs/1505.04597
    """

    def __init__(self, in_ch=3, out_ch=1):
        super().__init__()

        n1 = 64
        filters = [n1, n1 * 2, n1 * 4, n1 * 8, n1 * 16]

        self.Maxpool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.Maxpool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.Maxpool3 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.Maxpool4 = nn.MaxPool2d(kernel_size=2, stride=2)

        self.Conv1 = conv_block(in_ch, filters[0])
        self.Conv2 = conv_block(filters[0], filters[1])
        self.Conv3 = conv_block(filters[1], filters[2])
        self.Conv4 = conv_block(filters[2], filters[3])
        self.Conv5 = conv_block(filters[3], filters[4])

        # self.active = torch.nn.Sigmoid()
        self.interWin1 = nn.Sequential(
            ViTSwinTransformerBlock([256, 256], dim=filters[0], num_heads=1, window_size=8, shift_size=0),
            ViTSwinTransformerBlock([256, 256], dim=filters[0], num_heads=1, window_size=8, shift_size=4)
        )
        self.interWin2 = nn.Sequential(VitPatchMerging(filters[0]),
                                       ViTSwinTransformerBlock([128, 128], dim=filters[1], num_heads=2, window_size=8,
                                                               shift_size=0),
                                       ViTSwinTransformerBlock([128, 128], dim=filters[1], num_heads=2, window_size=8,
                                                               shift_size=4)
                                       )
        self.interWin3 = nn.Sequential(VitPatchMerging(filters[1]),
                                       ViTSwinTransformerBlock([64, 64], dim=filters[2], num_heads=4, window_size=8,
                                                               shift_size=0),
                                       ViTSwinTransformerBlock([64, 64], dim=filters[2], num_heads=4, window_size=8,
                                                               shift_size=4)
                                       )
        self.interWin4 = nn.Sequential(VitPatchMerging(filters[2]),
                                       ViTSwinTransformerBlock([32, 32], dim=filters[3], num_heads=8, window_size=8,
                                                               shift_size=0),
                                       ViTSwinTransformerBlock([32, 32], dim=filters[3], num_heads=8, window_size=8,
                                                               shift_size=4)
                                       )
        self.interWin5 = nn.Sequential(VitPatchMerging(filters[3]),
                                       ViTSwinTransformerBlock([16, 16], dim=filters[4], num_heads=16, window_size=8,
                                                               shift_size=0),
                                       ViTSwinTransformerBlock([16, 16], dim=filters[4], num_heads=16, window_size=8,
                                                               shift_size=4)
                                       )

        self.UNetDecoder = UnetDecoder(filters)

    def forward(self, input):
        e1 = self.Conv1(input)

        e2 = self.Maxpool1(e1)
        e2 = self.Conv2(e2)

        e3 = self.Maxpool2(e2)
        e3 = self.Conv3(e3)

        e4 = self.Maxpool3(e3)
        e4 = self.Conv4(e4)

        e5 = self.Maxpool4(e4)
        e5 = self.Conv5(e5)

        g1 = self.interWin1(e1)
        g2 = self.interWin2(g1)
        g3 = self.interWin3(g2)
        g4 = self.interWin4(g3)
        g5 = self.interWin5(g4)

        e1 = e1 + g1
        e2 = e2 + g2
        e3 = e3 + g3
        e4 = e4 + g4
        e5 = e5 + g5

        out = self.UNetDecoder([e1, e2, e3, e4, e5])

        return out


class SwinTransformer_UNetDecoder(nn.Module):
    """
    UNet - Basic Implementation
    Paper : https://arxiv.org/abs/1505.04597
    """

    def __init__(self, in_ch=3, out_ch=1):
        super().__init__()

        n1 = 64
        filters = [n1, n1 * 2, n1 * 4, n1 * 8, n1 * 16]

        self.Conv1 = conv_block(in_ch, filters[0])

        # self.active = torch.nn.Sigmoid()
        self.interWin1 = nn.Sequential(
            ViTSwinTransformerBlock([256, 256], dim=filters[0], num_heads=1, window_size=8, shift_size=0),
            ViTSwinTransformerBlock([256, 256], dim=filters[0], num_heads=1, window_size=8, shift_size=4)
        )
        self.interWin2 = nn.Sequential(VitPatchMerging(filters[0]),
                                       ViTSwinTransformerBlock([128, 128], dim=filters[1], num_heads=2, window_size=8,
                                                               shift_size=0),
                                       ViTSwinTransformerBlock([128, 128], dim=filters[1], num_heads=2, window_size=8,
                                                               shift_size=4)
                                       )
        self.interWin3 = nn.Sequential(VitPatchMerging(filters[1]),
                                       ViTSwinTransformerBlock([64, 64], dim=filters[2], num_heads=4, window_size=8,
                                                               shift_size=0),
                                       ViTSwinTransformerBlock([64, 64], dim=filters[2], num_heads=4, window_size=8,
                                                               shift_size=4)
                                       )
        self.interWin4 = nn.Sequential(VitPatchMerging(filters[2]),
                                       ViTSwinTransformerBlock([32, 32], dim=filters[3], num_heads=8, window_size=8,
                                                               shift_size=0),
                                       ViTSwinTransformerBlock([32, 32], dim=filters[3], num_heads=8, window_size=8,
                                                               shift_size=4)
                                       )
        self.interWin5 = nn.Sequential(VitPatchMerging(filters[3]),
                                       ViTSwinTransformerBlock([16, 16], dim=filters[4], num_heads=16, window_size=8,
                                                               shift_size=0),
                                       ViTSwinTransformerBlock([16, 16], dim=filters[4], num_heads=16, window_size=8,
                                                               shift_size=4)
                                       )

        self.UNetDecoder = UnetDecoder(filters)

    def forward(self, input):
        e1 = self.Conv1(input)

        g1 = self.interWin1(e1)
        g2 = self.interWin2(g1)
        g3 = self.interWin3(g2)
        g4 = self.interWin4(g3)
        g5 = self.interWin5(g4)

        out = self.UNetDecoder([g1, g2, g3, g4, g5])

        return out

if __name__ == '__main__':
    block = LGFO_PConv(3, 2).cuda()
    input = torch.rand(1, 3, 256, 256).cuda()
    output = block(input)
    print(input.size(), output.size())